package org.honeynet.hbbackend.cymruwhois;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.Socket;
import java.net.UnknownHostException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.text.ParseException;
import java.text.SimpleDateFormat;
import java.util.Arrays;
import java.util.EnumSet;
import java.util.regex.Pattern;

import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import javax.ejb.EJBException;
import javax.ejb.Schedule;
import javax.ejb.SessionContext;
import javax.ejb.Singleton;
import javax.ejb.Startup;
import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.jms.ConnectionFactory;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.Session;
import javax.jms.Topic;
import javax.jms.TopicSubscriber;
import javax.sql.DataSource;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Singleton
@Startup
public class CymruWhoisBinding {
	
	private final static String HOST = "v4.whois.cymru.com";
	private final static int PORT = 43;
	private final static int BATCH_SIZE = 3;
	private final static Pattern SPLIT_REGEX = Pattern.compile("\\s+\\|\\s+");
	private final static int RECV_TIMEOUT = 3000; 
	
	private static Logger log = LoggerFactory.getLogger(CymruWhoisBinding.class);
	
	@Resource(mappedName="jdbc/hbbackend")
	DataSource ds;
	
	@Resource(mappedName="jms/DurableConsumer/cymruwhois")
	private ConnectionFactory connectionFactory;
	
	@Resource(mappedName="jms/new_ip")
	private Topic dest;

	@Resource
	private SessionContext sctx;
	
	private boolean stop = false;
	
	
	@Schedule(second="*/5", minute="*", hour="*", persistent=false)
	@TransactionAttribute(TransactionAttributeType.REQUIRES_NEW)
	public void work() {
		log.trace("work() called");
		
		javax.jms.Connection mq = null;
		Session sess = null;
		TopicSubscriber sub = null;
		
		try {
			mq = connectionFactory.createConnection();
			sess = mq.createSession(true, 0);
			sub = sess.createDurableSubscriber(dest, "cymruwhois");
			log.trace("created durable subscription");
			log.trace("starting receive loop for {} messages", BATCH_SIZE);
			mq.start();
			
			String ips[] = new String[BATCH_SIZE];
			int i = 0;
			while (i < BATCH_SIZE) {
				if (stop) {
					log.trace("encountered stop request: returning");
					return;
				}
				
				Message jmsMsg;
				do {
					if (stop) {
						log.trace("encountered stop request: returning");
						sctx.setRollbackOnly();
						return;
					}
					log.trace("receive");
					jmsMsg = sub.receive(RECV_TIMEOUT);
				} while (jmsMsg == null);
				
				//TODO message validation as the service is ipv4 only
				ips[i] = jmsMsg.getStringProperty("ip");
				
				i++;
			}
			
			store(retrieve(ips));
			return;
			
		} catch (JMSException e) {
			log.error("got JMSException", e);
			throw new EJBException(e);
		} finally {
			try {
//				if (sub != null) sub.close();
//				if (sess != null) sess.close();
				if (mq != null) mq.close();
			} catch (JMSException e) {
				log.error("got JMSException", e);
				throw new EJBException(e);
			}
		}
		
	}
	
	
	private Result[] retrieve(String[] ips) {
		log.trace("retrieve() called");
		
		String emsg;
		
		Socket socket = null;
		BufferedReader in = null;
		PrintWriter out = null;
		

		try {
			log.trace("connecting to {}:{}", HOST, PORT);
			socket = new Socket(HOST, PORT);
			
			in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
			out = new PrintWriter(socket.getOutputStream(), true);

			String response = null;

			log.trace("beginning request");
			out.print("begin\n");
			out.flush();

			response = in.readLine();
			if (response.matches("^Bulk mode; whois.cymru.com.*")) {
				log.trace("response: {}", response);
			} else {
				emsg = "couldn't parse first line of response: " + response;
				log.error(emsg);
				throw new EJBException(emsg);
			}

			out.print("verbose\n");
			out.flush();

			Result[] results = new Result[ips.length];
			int i = 0;

			for (String ip : ips) {
				log.trace("sending ip = {}", ip);
				out.print(ip);
				out.print("\n");
				out.flush();

				log.trace("waiting for result for ip = {}", ip);
				response = in.readLine(); // TODO read timeout?
				log.trace("response line: {}", response);

				String[] responseFields = SPLIT_REGEX.split(response);
				log.trace("parsed response fields = {}", Arrays.asList(responseFields));
				
				// general check
				if (responseFields.length != Fields.values().length) {
					emsg = "invalid number of fields in response";
					log.error(emsg);
					throw new EJBException(emsg);
				}
				
				// check ip
				String ipInResponse = responseFields[Fields.ip.ordinal()].trim();
				if (ipInResponse.equals(ip) == false) {
					emsg = "ip in response does not match requested ip: got = " + ipInResponse + " requested = " + ipInResponse;
					log.error(emsg);
					throw new EJBException(emsg);
				}

				// general validation
				for (Fields field : EnumSet.allOf(Fields.class)) {
					String value = responseFields[field.ordinal()].trim();
					log.trace("read field: {} = {}", field.name(), value);
					if (field.validate(value) == false) {
						emsg = "field does not validate: " + field.name();
						log.error(emsg);
						throw new EJBException(emsg);
					}
				}
				
				
				Result res = new Result();
				
				res.ip = ip;
				res.asn = Long.parseLong(responseFields[Fields.asn.ordinal()]);

				java.util.Date date = null;
				try {
					date = new SimpleDateFormat("yyyy-MM-dd")
							.parse(responseFields[Fields.allocated.ordinal()]);
				} catch (ParseException e) {
					emsg = "couldn't parse date: " + Fields.allocated.name() + " = " +responseFields[Fields.allocated.ordinal()];
					log.error(emsg, e);
					throw new EJBException(emsg);
				}
				res.allocated = new java.sql.Date(date.getTime());
				res.as_name = responseFields[Fields.as_name.ordinal()];
				res.cc = responseFields[Fields.cc.ordinal()];
				res.registry = responseFields[Fields.registry.ordinal()];
				res.bgp_prefix = responseFields[Fields.bgp_prefix.ordinal()];

				results[i] = res;
				i++;
			}
			
			log.trace("ending request");
			out.print("end\n");

//			out.close();
//			in.close();
//			socket.close();
			
			for (Result res : results) {
				log.trace("result: {}", res);
			}
			
			return results;
			
		} catch (UnknownHostException e) {
			log.error("got UnknownHostException", e);
			throw new EJBException(e);
		} catch (IOException e) {
			log.error("got IOException", e);
			throw new EJBException(e);
		} finally {
			try {
				if (out != null) out.close();
				if (in != null)	in.close();
				if (socket != null)	socket.close();
			} catch (IOException e) {
				log.error("got IOException", e);
				throw new EJBException(e);
			}
		}
	}
	
	
	private void store(Result[] results) {
		log.trace(" store() called");
		
		Connection db = null;
		PreparedStatement pStmt;
		
		try {
			db = ds.getConnection();
			pStmt = db.prepareStatement("insert into asn_cymru (ip, asn, as_name, cc, registry, allocated, bgp_prefix) values (inet(?),?,?,?,?,?,inet(?))");
			
			for (Result res : results) {
				log.trace("inserting into asn_cymru (ip = {})", res.ip);

				pStmt.setString(1, res.ip);
				pStmt.setLong(2, res.asn);
				pStmt.setString(3, res.as_name);
				pStmt.setString(4, res.cc);
				pStmt.setString(5, res.registry);
				pStmt.setDate(6, res.allocated);
				pStmt.setString(7, res.bgp_prefix);
				
				pStmt.executeUpdate();
			}
			
			pStmt.close();
			
		} catch (SQLException e) {
			log.error("got SQLException", e);
			throw new EJBException(e);
		}
		finally {
			try {
				if (db != null)	db.close();
			} catch (SQLException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
		}
		

	}
	
	
	@PreDestroy
	public void stop() {
		stop = true;
	}
	
	
	private static enum Fields {
		asn,
		ip,
		bgp_prefix,
		cc,
		registry,
		allocated,
		as_name;
		
		public boolean validate(String value) {
			return true;
		}
	}
	
	
	private static class Result {
		String ip;
		
		long asn;
		String as_name;
		String cc;
		String registry;
		java.sql.Date allocated;
		String bgp_prefix;
		
		public String toString() {
			return
				Fields.ip.name() + " = " + ip + " | " +
				Fields.asn.name() + " = " + Long.toString(asn) + " | " +
				Fields.as_name.name() + " = " + as_name + " | " +
				Fields.cc.name() + " = " + cc + " | " +
				Fields.registry.name() + " = " + registry + " | " +
				Fields.allocated.name() + " = " + allocated + " | " +
				Fields.bgp_prefix.name() + " = " + bgp_prefix;
		}
	}
	
}
